#! /usr/bin/env python3
import numpy as np
from PIL import Image, ImageFile
from flask import Flask, jsonify, request, render_template
from keras.models import load_model

ImageFile.LOAD_TRUNCATED_IMAGES = True

from tqdm import tqdm

import os, pickle, uuid
from glob import glob
from functools import reduce
import time, datetime
import logging

from libs.models import AreaDetect
from libs.tools import calcAreaBoxs, calcSourceBoxs, calcAccurateBoxs, textCleaner, textDeskew

from libs.recall import MNBClassifier, formatInput

app = Flask(__name__)

file_handler = logging.FileHandler('logs.txt', encoding='utf-8')
app.logger.addHandler(file_handler)

import requests
import json
with open('config.json', 'r', encoding='utf-8') as f:
    config = json.load(f)

document_clf = load_model(config['DOCUMENT_CLF']['MODEL'], compile=False)
layout_clf = load_model(config['LAYOUT_CLF']['MODEL'], compile=False)
# type_clf = load_model(config['TYPE_CLF']['MODEL'], compile=False)

clf_map = {
    'document': document_clf, 'document_classifier': document_clf,
    'layout': layout_clf, 'layout_classifier': layout_clf,
    # 'type': type_clf, 'type_classifier': type_clf,
}

deskew_model = load_model(config['TEXTDESKEW_MODEL'], compile=False)

def load_labels(path):
    with open(path, 'rb') as f:
        labels = pickle.load(f)
    return labels

document_labels = load_labels(config['DOCUMENT_CLF']['LABELS'])
layout_labels = load_labels(config['LAYOUT_CLF']['LABELS'])
# type_labels = load_labels(config['TYPE_CLF']['LABELS'])

labels_map = {
    'document': document_labels, 'document_classifier': document_labels,
    'layout': layout_labels, 'layout_classifier': layout_labels,
    # 'type': type_labels, 'type_classifier': type_labels,
}

model_paths = glob(os.path.join(config['AREA_DETECT_MODELS_DIR'], '*.h5'))
models = {}
for model_path in tqdm(model_paths, ncols=80):
    model_name = os.path.basename(model_path).split('.')[0]
    model_data = load_model(model_path, compile=False)
    models[model_name] = model_data

recall_paths = glob(os.path.join(config['RECALL_MODELS_DIR'], '*.pkl'))
recall_paths += glob(os.path.join(config['RECALL_MODELS_DIR'], '*.h5'))
recall_models = {}
for recall_path in tqdm(recall_paths, ncols=80):
    recall_name = os.path.basename(recall_path).split('.')[0]
    if os.path.splitext(recall_path)[-1].lower() == '.pkl':
        recall_data = MNBClassifier()
        recall_data.load(recall_path)
    else:
        recall_data = load_model(recall_path, compile=False)
    recall_models[recall_name] = recall_data

with open(config['CHI_SIM'], encoding='utf-8') as f:
    chi_sim = [w[0] for w in f.readlines()[:3500]]


def classifierImage(input_img, clf, labels=None):
    # Test code.
    input_shape = clf.get_input_shape_at(0)
    if isinstance(input_img, str):
        img = Image.open(input_img)
    elif isinstance(input_img, Image.Image):
        img = input_img
    else:
        raise TypeError('classifierImage input_img is path string or pillow image object')

    if labels is None:
        labels = list(range(clf.get_output_shape_at(0)[-1]))

    img_arr_formated = AreaDetect.formatImage(img, input_shape[1:])

    X = img_arr_formated.astype('float32')
    X = 1 - X / 255
    X = X.reshape(-1, input_shape[1], input_shape[2], input_shape[3])

    frame = {}
    probs = clf.predict(X)[0]
    predict = (probs == probs.max()).dot(np.arange(probs.size))
    probs = dict(zip(labels, probs.tolist()))
    frame['Predict'] = labels[int(predict)]
    frame['Recommend'] = sorted(probs.items(), key=lambda item: item[1], reverse=True)[:config['RECOMMEND_TYPE_COUNT']]

    return frame

def classifierImages(directory, clf, labels):
    filename_list = glob(os.path.join(directory, '*'))

    frame = []
    for filename in filename_list:
        row = classifierImage(filename, clf, labels)
        row['Path'] = filename
        frame.append(row)
    return frame

def detectImage(input_img, model, labels=None, T=False):
    # Test code.
    input_shape = model.get_input_shape_at(0)
    if isinstance(input_img, str):
        img = Image.open(input_img)
    elif isinstance(input_img, Image.Image):
        img = input_img
    else:
        raise TypeError('detectImage input_img is path string or pillow image object')

    img_arr_formated = AreaDetect.formatImage(img, input_shape[1:])
    X = img_arr_formated.astype('float32')
    X = 1 - X / 255
    X = X.reshape(-1, input_shape[1], input_shape[2], input_shape[3])

    predict_masks = model.predict(X)[0]

    a = max(input_shape[1:-1])
    classes_num = predict_masks.shape[-1]
    boxs_list = []
    for i in range(classes_num):
        boxs_p = calcAreaBoxs(predict_masks[:,:, i])
        boxs_p_arr = calcSourceBoxs(img, boxs_p, a)
        if T:
            boxs_list.append(boxs_p_arr.tolist())
        else:
            boxs_p_arr_ac = calcAccurateBoxs(np.array(img.convert('L')), boxs_p_arr)
            row = {'L1': boxs_p_arr.tolist(), 'L2': boxs_p_arr_ac.tolist()}
            boxs_list.append(row)
    
    frame = {}
    if labels is None:
        if isinstance(boxs_list[0], dict):
            frame['L1'] = reduce(lambda a, b: a + b,
                                 [item['L1'] for item in boxs_list])
            frame['L2'] = reduce(lambda a, b: a + b,
                                 [item['L2'] for item in boxs_list])
        else:
            frame = reduce(lambda a, b: a + b, boxs_list)
    else:
        for i, label in enumerate(labels):
            frame[label] = boxs_list[i]

    if T: return img, frame
    return frame

def detectImages(directory, model):
    filename_list = glob(os.path.join(directory, '*'))

    frame = []
    for filename in filename_list:
        filename = os.path.abspath(filename)
        row = [filename, detectImage(filename, model)]
        frame.append(row)
    with open('fbs.pkl', 'wb') as f:
        pickle.dump(frame, f)

class FullPipeline:
    def __init__(self, path):
        self.path = path

    def run(self):
        frame = None
        if os.path.isfile(self.path):
            frame = self.fullPipeline(self.path)
        elif os.path.isdir(self.path):
            frame = self.fullPipelines(self.path)
        return frame

    def fullPipelines(self, directory):
        input_paths = glob(os.path.join(directory, '*'))

        frame = []
        for input_path in input_paths:
            row = self.fullPipeline(input_path)
            frame.append(row)
        return frame

    def fullPipeline(self, input_path):
        frame = {'Path': input_path}
        filename = input_path

        self.input_basename = os.path.basename(filename)

        # >>> Test code 03-17

        # input_img = Image.open(input_path)
        # frame['Type'] = classifierImage(input_img, type_clf, type_labels)
        # if frame['Type']['Predict'] == config['AREA_DETECT_BRANCH']['TEXT_IMAGE']:
        #     frame = self.textBranch(frame, input_path, input_img)
        # else:
        #     frame = self.photoBranch(frame, input_img)
        # return frame

        input_img = Image.open(input_path)
        frame['Document'] = classifierImage(input_img, document_clf, document_labels)
        if frame['Document']['Predict'] == config['AREA_DETECT_BRANCH']['NON_TEXT_IMAGE']:
            frame = self.photoBranch(frame, input_img)
        else:
            frame = self.textBranch(frame, input_path, input_img)
        return frame

        # <<<

    def photoBranch(self, frame, input_img):
        frame['IsText'] = False
        
        # Test code.
        frame['Boxs'] = []
        if '非文本图片' not in config['AREA_DETECT_MODEL_DICT']: return frame

        model_name = config['AREA_DETECT_MODEL_DICT']['非文本图片']
        if model_name in models.keys() and models[model_name] is not None:
            model = models[model_name]
            frame['Boxs'] = detectImage(input_img, model) 

        return frame

    @staticmethod
    def addExtension(path):
        exts = ['jpg', 'jpeg', 'JPG', 'JPEG']
        if os.path.splitext(path)[-1] in exts:
            return path
        return os.path.splitext(path)[0] + '.jpg'

    def textBranch(self, frame, input_path, input_img):
        frame['IsText'] = True
        if config['USE_DESKEW']:
            basename = os.path.basename(input_path)
            basename = '-{}'.format(uuid.uuid4()).join(os.path.splitext(basename))
            output_path = os.path.join(config['TMP_DIR'], basename)
            output_path = self.addExtension(output_path)
            frame['PathDeskew'] = output_path
            deskew = textDeskew(input_path, output_path, deskew_model,
                                limit=config['DESKEW_LIMIT'], no_bias=config['DESKEW_NO_BIAS'])
            input_img = Image.fromarray(deskew)

        # >>> Test code 03-17

        # frame['Document'] = classifierImage(input_img, document_clf, document_labels)

        # frame = self.recallBranch(frame, input_img)

        # if frame['Document']['Predict'] == config['AREA_DETECT_BRANCH']['NON_DOCUMENT']:
        #     frame = self.layoutBranch(frame, input_img)
        # else:
        #     frame = self.documentBranch(frame, input_img)
        # if config['USE_CLEANER']:
        #     basename = os.path.basename(input_path)
        #     basename = '-{}'.format(uuid.uuid4()).join(os.path.splitext(basename))
        #     output_path = os.path.join(config['TMP_DIR'], basename)
        #     output_path = self.addExtension(output_path)
        #     frame['PathCleaner'] = output_path
        #     textCleaner(input_img, output_path)
        # return frame

        frame = self.recallBranch(frame, input_img)
        if frame['Document']['Predict'] == config['AREA_DETECT_BRANCH']['NON_DOCUMENT']:
            frame = self.layoutBranch(frame, input_img)
        else:
            frame = self.documentBranch(frame, input_img)
        if config['USE_CLEANER']:
            basename = os.path.basename(input_path)
            basename = '-{}'.format(uuid.uuid4()).join(os.path.splitext(basename))
            output_path = os.path.join(config['TMP_DIR'], basename)
            output_path = self.addExtension(output_path)
            frame['PathCleaner'] = output_path
            textCleaner(input_img, output_path)
        return frame

        # <<<

    @staticmethod
    def image_to_string(path):
        url = 'http://19z8857a58.iok.la:88/archives/capi/ocr.do'
        data = {'userImageFileName': '0123456789_kk.jpg'}
        files = {'file': open(path, 'rb')}
        res = requests.post(url, data=data, files=files)
        txt = ''
        try:
            txt = json.loads(res.text).get('result', '')
        except: pass
        return txt

    def recallBranch(self, frame, input_img):
        document_predict = frame['Document']['Predict']
        if document_predict in config['RECALL']:
            recall_name = config['RECALL'][document_predict]
            if recall_name in recall_models.keys():
                recall_clf = recall_models[recall_name]
                if isinstance(recall_clf, MNBClassifier):
                    txt = self.image_to_string(frame['Path'])
                    if not int(recall_clf.predict(formatInput(txt, chi_sim))):
                        frame['Document']['Predict'] = config['AREA_DETECT_BRANCH']['NON_DOCUMENT']
                else:
                    if not int(classifierImage(input_img, recall_clf)['Predict']):
                        frame['Document']['Predict'] = config['AREA_DETECT_BRANCH']['NON_DOCUMENT']
        return frame

    def layoutBranch(self, frame, input_img):
        frame['IsDocument'] = False
        frame['Layout'] = classifierImage(input_img, layout_clf, layout_labels)
        if frame['Layout']['Predict'] in config['AREA_DETECT_BRANCH']['FRONT_COVER'].split(','):
            frame['IsLayout'] = True
            model_name = config['AREA_DETECT_MODEL_DICT'][config['AREA_DETECT_BRANCH']['FRONT_COVER']]
            if model_name in models.keys() and models[model_name] is not None:
                model = models[model_name]
                frame['Boxs'] = detectImage(input_img, model) 
            else:
                frame['Boxs'] = []
        else:
            frame['IsLayout'] = False
        return frame

    def documentBranch(self, frame, input_img):
        # Test code.
        if config['EXT_DOCUMENT_LAYOUT_BRANCH']:
            frame = self.layoutBranch(frame, input_img)
            frame['LayoutBoxs'] = frame.get('Boxs', []) # frame['Boxs']

        frame['IsDocument'] = True
        frame['Boxs'] = []
        predict_doc = frame['Document']['Predict']
        if (predict_doc in config['AREA_DETECT_SLICE_MODEL_DICT'] and
            config['AREA_DETECT_SLICE_MODEL_DICT'][predict_doc] is not None):
            frame = self.sliceImage(frame, input_img)
        elif predict_doc in config['AREA_DETECT_MODEL_DICT']:
            frame['Boxs'] = self.detectDOC(frame, input_img)
        return frame

    def sliceImage(self, frame, input_img):
        predict_doc = frame['Document']['Predict']
        model_name = config['AREA_DETECT_SLICE_MODEL_DICT'][predict_doc]
        if model_name in models.keys() and models[model_name] is not None:
            model = models[model_name]
            img, boxs = detectImage(input_img, model, T=True)

            # Test code.
            basename = '{}.jpg'.format(uuid.uuid4())
            # basename = self.input_basename

            # Test code.
            W, H = input_img.size

            for i, ([y0, x0], [y1, x1]) in enumerate(boxs):
                a = 0 - x0 if x0 < 0 else 0
                b = 0 - y0 if y0 < 0 else 0
                
                w, h = x1 - x0, y1 - y0
                wr, hr = w / W, h / H

                if (min(wr, hr) < config['FILTER_SIZE'][0] or
                    max(wr, hr) < config['FILTER_SIZE'][1]): continue

                img_crop = img.crop((x0+a, y0+b, x1+a, y1+b))
                slice_filename = '-{}'.format(i).join(os.path.splitext(basename))
                slice_filename = os.path.join(config['TMP_DIR'], slice_filename)

                # Test code.
                slice_filename = self.addExtension(slice_filename)

                # Test code.
                if config['USE_SLICE_DESKEW']:
                    deskewed_arr = textDeskew(img_crop, slice_filename, deskew_model, no_bias=True)
                    deskewed_img = Image.fromarray(deskewed_arr)
                else:
                    deskewed_img = img_crop
                    deskewed_img.save(slice_filename)

                if predict_doc in config['AREA_DETECT_MODEL_DICT']:
                    slice_boxs = self.detectDOC(frame, deskewed_img)
                    frame['Boxs'].append({slice_filename: slice_boxs})
        return frame

    def detectDOC(self, frame, input_img):
        predict_doc = frame['Document']['Predict']
        model_name = config['AREA_DETECT_MODEL_DICT'][predict_doc]
        boxs = []
        if model_name in models.keys() and models[model_name] is not None:
            model = models[model_name]
            if predict_doc in config['AREA_DETECT_LABELS_DICT']:
                labels = config['AREA_DETECT_LABELS_DICT'][predict_doc]
                boxs = detectImage(input_img, model, labels=labels) 
            else:
                boxs = detectImage(input_img, model) 
        return boxs

@app.route('/classifier/<clf_name>/<token>/<path:path>')
def classifier(clf_name, token, path):
    try:
        if token != config['TOKEN']:
            return 'Please input your token text!'
        if clf_name in clf_map.keys():
            clf = clf_map[clf_name]
            labels = labels_map[clf_name]
        else:
            return 'No find {}'.format(clf_name)

        frame = None
        full_path = os.path.join(config['ROOT_PATH'], path)
        if os.path.isfile(full_path):
            frame = classifierImage(full_path, clf, labels)
        elif os.path.isdir(full_path):
            frame = classifierImages(full_path, clf, labels)

        return jsonify(frame)
    except BaseException as e:
        app.logger.error('Find error in classifier %r', e)

@app.route('/detect_area/<model_name>/<token>/<path:path>')
def detectArea(model_name, token, path):
    try:
        if token != config['TOKEN']:
            return 'Please input your token text!'
        if model_name in models.keys():
            model = models[model_name]
        else:
            return 'No find {}'.format(model_name)

        full_path = os.path.join(config['ROOT_PATH'], path)
        if os.path.isfile(full_path):
            boxs = detectImage(full_path, model)
            return jsonify(boxs)
        elif os.path.isdir(full_path):
            detectImages(full_path, model)
            return 'Detect finish'
    except BaseException as e:
        app.logger.error('Find error in detectArea %r', e)

@app.route('/full/<token>/<path:path>')
def fullDetect(token, path):
    # try:
        if token != config['TOKEN']:
            return 'Please input your token text!'

        full_path = os.path.join(config['ROOT_PATH'], path)
        fullPipline = FullPipeline(full_path)
        frame = fullPipline.run()

        return jsonify(frame)
    # except BaseException as e:
    #     app.logger.error('Find error in fullDetect %r', e)

@app.route('/update_classifier/<model_name>/<token>')
def updateClassifier(model_name, token):
    global document_clf, layout_clf, document_labels, layout_labels, deskew_model

    try:
        if token != config['TOKEN']:
            return 'Please input your token text!'

        if model_name == 'document_clf':
            document_clf = load_model(config['DOCUMENT_CLF']['MODEL'], compile=False)
            document_labels = load_labels(config['DOCUMENT_CLF']['LABELS'])
        elif model_name == 'layout_clf':
            layout_clf = load_model(config['LAYOUT_CLF']['MODEL'], compile=False)
            layout_labels = load_labels(config['LAYOUT_CLF']['LABELS'])
        elif model_name == 'deskew_model':
            deskew_model = load_model(config['TEXTDESKEW_MODEL'], compile=False)
        else: return 'fail'

        return 'success'
    except BaseException as e:
        app.logger.error('Find error in updateClassifier %r', e)

@app.route('/update_areadetect/<model_name>/<token>')
def updateAreaDetect(model_name, token):
    global models

    try:
        if token != config['TOKEN']:
            return 'Please input your token text!'

        if (model_name in [os.path.splitext(filename)[0]
                           for filename in os.listdir(config['AREA_DETECT_MODELS_DIR'])]):
            model_path = os.path.join(config['AREA_DETECT_MODELS_DIR'], model_name+'.h5')
            models[model_name] = load_model(model_path, compile=False)
        else: return 'fail'

        return 'success'
    except BaseException as e:
        app.logger.error('Find error in updateAreaDetect %r', e)

@app.route('/update_recall/<model_name>/<token>')
def updateRecall(model_name, token):
    global recall_models

    try:
        if token != config['TOKEN']:
            return 'Please input your token text!'

        if (model_name in [os.path.splitext(filename)[0]
                             for filename in os.listdir(config['RECALL_MODELS_DIR'])]):
            model_path = os.path.join(config['RECALL_MODELS_DIR'], model_name+'.h5')
            if os.path.isfile(model_path):
                recall_models[model_name] = load_model(model_path, compile=False)
            else:
                model_path = os.path.join(config['RECALL_MODELS_DIR'], model_name+'.pkl')
                recall_models[model_name] = MNBClassifier()
                recall_models[model_name].load(model_path)
        else: return 'fail'

        return 'success'
    except BaseException as e:
        app.logger.error('Find error in updateRecall %r', e)

@app.route('/config/<token>', methods=['GET', 'POST'])
def setConfig(token):
    global config

    try:
        if token != config['TOKEN']:
            return 'Please input your token text!'

        if request.method == 'POST':
            new_config = request.get_json()
            if new_config is None:
                return jsonify({'error': 'invalid payload'})
            config.update(new_config)
            with open('./config.json', 'w') as f:
                json.dump(config, f, ensure_ascii=False, indent=4, sort_keys=True)
            return jsonify(config)

        return render_template('config.html', Args=dict(config_str=json.dumps(config), token=config['TOKEN']))
    except BaseException as e:
        app.logger.error('Find error in setConfig %r', e)

monitorDict = {}
@app.route('/monitor/<token>', methods=['GET', 'POST'])
def trainingMonitor(token):
    global monitorDict

    try:
        if token != config['TOKEN']:
            return 'Please input your token text!'
        if request.method == 'POST':
            if request.form.get('model_filename') is None:
                return jsonify({'warning': 'the request not from train shell'})
            model_filename = request.form.get('model_filename')
            epochs_num = int(request.form.get('epochs_num'))
            keys = ['epoch', 'loss', 'metric', 'est_time']
            epoch, loss, metric, est_time = [request.form.get(k) for k in keys]
            epoch, loss, metric = int(epoch), float(loss), float(metric)
            data = {
                'epochs_num': epochs_num,
                'epoch': epoch, 'loss': loss, 'metric': metric,
                'process_rate': int(epoch / epochs_num * 100),
                'est_time': est_time}

            if model_filename in monitorDict.keys():
                monitorDict[model_filename].update(data)
            else:
                monitorDict[model_filename] = data

            now_datetime = datetime.datetime.today()
            monitorDict[model_filename]['time'] = time.time()
            monitorDict[model_filename]['now_time'] = now_datetime.strftime('%H:%M:%S')
            return jsonify(monitorDict)

        status_dict_timeout = 60
        current_time = time.time()
        monitorDict = {k: v for k, v in monitorDict.items() if current_time - v['time'] < status_dict_timeout}
        return render_template('monitor.html', Args=monitorDict)
    except BaseException as e:
        app.logger.error('Find error in trainingMonitor %r', e)

@app.route('/batch/<token>')
def batchUpdate(token):
    try:
        if token != config['TOKEN']:
            return 'Please input your token text!'
        return render_template('batch.html', Args=dict(config_str=json.dumps(config), token=config['TOKEN']))
    except BaseException as e:
        app.logger.error('Find error in batch %r',e)


@app.route('/upAreaDectMd/<model_type>/<token>')
def areadetectmodel(model_type, token):
    try:
        if token != config['TOKEN']:
            return jsonify({'state': False, 'msg': 'Please input your token text!'})
        if model_type is None or len(model_type) == 0:
            return jsonify({'state': False, 'msg': '证件类型不能为Null'})

        model_types = model_type.split(',')
        for mt in model_types:
            config['AREA_DETECT_MODEL_DICT'][mt] = 'cover'

        with open('config.json', 'w', encoding='utf-8') as outfile:
            outfile.write(json.dumps(config, sort_keys=False, indent=2, separators=(',', ':'), ensure_ascii=False))
            outfile.close()

        return jsonify({'state': True})
    except BaseException as e:
        app.logger.error('Find error in areadetectmodel %r', e)
        return jsonify({'state': False, 'msg': '系统异常!'})

if __name__ == '__main__':
    app.run(debug=False, port=config['PORT'])
